# -*- coding: utf-8 -*-
from IMAGE_QUERY.extract_cnn_vgg16_keras import VGGNet
import pandas as pd
import numpy as np
from numpy import linalg as LA
from pymongo import MongoClient


def query(query_path, filename, query_num):
    if query_num > 50:
        query_num = 50

    # read in indexed images' feature vectors and corresponding image names
    client = MongoClient('localhost', 27017)
    db = client.toutiao
    collection = db.toutiao_feats_04

    items = collection.find({})
    df_data = pd.DataFrame(list(items))
    df_new = df_data.drop_duplicates(subset='img_name', keep='first', inplace=False)

    feats = np.array([i/LA.norm(i) for i in df_new.feats])
    img_names = df_new.img_name.tolist()

    print("--------------------------------------------------")
    print("searching starts......")
    print("--------------------------------------------------")

    model = VGGNet()
    query_vec = model.extract_feat(query_path)
    query_vec_norm = query_vec/LA.norm(query_vec)

    temp_dict = {'img_name': filename, 'feats': query_vec.tolist()}
    print(filename)
    print(img_names[-5:])
    if filename not in img_names:
        collection.insert_one(temp_dict)

    scores = np.dot(query_vec_norm, feats.T)
    rank_ID = np.argsort(scores)[::-1]

    query_num = int(query_num)
    imlist = [img_names[index] for i, index in enumerate(rank_ID[0:query_num])]
    client.close()
    return imlist


if __name__ == '__main__':

    query_path = 'D:/data_set/image_test/Big_Hero_6.jpg'
    filename = 'Big_Hero_6.jpg'
    query_num = 10
    output = query(query_path, filename, query_num)
    print(output)



